options(scipen=999)
library(AlphaMME)

args = commandArgs(trailingOnly=TRUE)
h2 = as.numeric(args[1])
doTraining = args[2]
folder = args[3]


totalGeneticVariance = h2/(1-h2)


# folder = "test"
# doTraining = "train"
# nEffects = 200
# totalGeneticVariance = 1

#Load in base population data
print("Reading genotypes")

genotypes = as.matrix(read.table("genotypes.children"))

print("Reading finished")

ids = genotypes[,1]
genotypes = genotypes[,-1]

### Create baseline effects to use with no training.
nEffects = as.numeric(ncol(genotypes))

effectSize = totalGeneticVariance
if(effectSize != 0){
    effectsTrue = rnorm(nEffects, mean = 0, sd = 1)
} else{
    effectsTrue = rep(0, nEffects)
}
nLoci = ncol(genotypes)
loci = sample(1:nLoci, size = nEffects)

effectsExpanded = rep(0, nLoci)
effectsExpanded[loci] = effectsTrue
# Generate phenotypes + breeding values
print("Generating BV")

bv_pop = genotypes %*% effectsExpanded 

# Normalize so this has variance equal to genetic variance.
scale = sqrt(effectSize)/sd(bv_pop)

bv_pop = bv_pop * scale
effectsExpanded = effectsExpanded * scale

pheno_pop = bv_pop + rnorm(length(bv_pop), 0, 1)
print("Generating BV finished")


print(var(bv_pop))
print(var(pheno_pop))

write.table(effectsExpanded, paste0(folder, "/qtl.0"))
write.table(c(beta  = 0, Ve = 1), paste0(folder, "/params.", "0"))

bv = cbind(ids, bv_pop)
write.table(bv, file.path(folder, "bv.txt"), row.names=F, col.names=F, quote=F)

pheno = cbind(ids, pheno_pop)
write.table(pheno, file.path(folder, "pheno.txt"), row.names=F, col.names=F, quote=F)

h = var(bv[,2])/var(pheno[,2])
write.table(h, file.path(folder, "h2.txt"), row.names=F, col.names=F, quote=F)

# Generate values for other traits.
print("IO completed")

if(doTraining == "train"){
    genotypes = as.matrix(read.table("training_genotypes.txt"))[,-1]


    bv_pop = genotypes[1:1000,] %*% effectsExpanded 
    pheno_pop = bv_pop + rnorm(nrow(genotypes), 0, 1)

    # for(trainingSize in c(1000, 2000, 4000, 6000, 8000, 10000)){
    for(trainingSize in c(100, 250, 500, 750, 1000)){
        ### Get estimated effect sizes.
        y = matrix(pheno_pop[1:trainingSize])
        X = matrix(rep(1, trainingSize))
        M = genotypes[1:trainingSize,]

        vals = AlphaMME::solveRRBLUP(y, X, M)

        effectsInfer = vals$u

        write.table(effectsInfer, paste0(folder, "/qtl.", trainingSize))
        write.table(c(beta  = vals$beta, Ve = vals$Ve), paste0(folder, "/params.", trainingSize))
    }
}



